package em;

import java.util.ArrayList;

import main_app.MainApp;

/**
 * Class represents the distributions parameters (theta, as seen at class)
 */
public class Parameters {

	private double[] topicsDistribution; // probability for each topic in corpus
	private ArrayList<ArrayList<Double>> wordByTopicDistribution; // probability for each word given each topic
	private int numWords; // number of different words in corpus
	
	/**
	 * Constructor- creates new parameters (theta)
	 * @param numWords number of different words in corpus
	 */
	public Parameters(int numWords) {
		
		this.numWords = numWords;
		
		// initialize distribution over topics
		topicsDistribution = new double[MainApp.NUM_TOPICS];
		
		// initialize distribution over words for each topic
		wordByTopicDistribution = new ArrayList<ArrayList<Double>>();
		for (int topicIndex = 0 ; topicIndex < MainApp.NUM_TOPICS ; ++topicIndex) {
			
			// initialize distribution over words for current topic
			ArrayList<Double> currTopicDist = new ArrayList<Double>(numWords);
			for (int wordIndex = 0 ; wordIndex < numWords ; ++wordIndex) {
				currTopicDist.add(0.0);
			}
			
			wordByTopicDistribution.add(currTopicDist);
		}
	}
	
	/**
	 * Returns probability in corpus of indicated topic: P(x)
	 * @param topicIndex index of topic to check
	 * @return probability of topic matching given index
	 */
	public double getTopicProb(int topicIndex) {
		
		// check boundaries
		if (topicIndex < 0 || topicIndex > MainApp.NUM_TOPICS) {
			return 0;
		}
		
		// return probability of topic matching given index
		return topicsDistribution[topicIndex];
	}
	
	/**
	 * Returns conditional probability in corpus of indicated word given 
	 * indicated topic: P(w|x)
	 * @param wordIndex index of word to check
	 * @param topicIndex index of topic to check
	 * @return conditional probability of indicated word given indicated topic
	 */
	public double getWordProb(int wordIndex, int topicIndex) {
		
		// check boundaries
		if (topicIndex < 0 || topicIndex >= MainApp.NUM_TOPICS || wordIndex < 0
				|| wordIndex >= numWords) {
			return 0;
		}
		
		// return probability for word given topic
		return wordByTopicDistribution.get(topicIndex).get(wordIndex);
	}
	
	/**
	 * Sets probability in corpus of indicated topic.
	 * If topic index is invalid, does not set any value 
	 * @param topicIndex index of topic
	 * @param prob new probability value
	 */
	public void setTopicProb(int topicIndex, double prob) {
		
		// check boundaries
		if (topicIndex < 0 || topicIndex > MainApp.NUM_TOPICS) {
			return;
		}
		
		// set probability
		topicsDistribution[topicIndex] = prob;
	}
	
	/**
	 * Sets conditional probability in corpus of indicated word given
	 * indicated topic.
	 * If topic or word index is invalid, does not set any value
	 * @param wordIndex index of word
	 * @param topicIndex index of topic
	 * @param prob
	 */
	public void setWordByTopicProb(int wordIndex, int topicIndex, double prob) {

		// check boundaries
		if (topicIndex < 0 || topicIndex >= MainApp.NUM_TOPICS || wordIndex < 0
				|| wordIndex >= numWords) {
			return;
		}
		
		// set probability for word given topic
		wordByTopicDistribution.get(topicIndex).set(wordIndex, prob);
	}
}
